#pragma once
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vector>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <geometry_msgs/msg/point_stamped.hpp>
#include <tf2_ros/buffer.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#include <rclcpp/rclcpp.hpp>
#include "yolov8.h"

#include "image_drawing.h"
#define OFFSET 0.20
// 目标消息类型
struct DetectedObject {
    std::string label;
    float confidence;
    cv::Rect bbox;
    geometry_msgs::msg::PointStamped position_3d;
};

// 定义目标检测结果列表
struct DetectionResult {
    std_msgs::msg::Header header;
    cv::Mat processed_image;
    std::vector<DetectedObject> detected_objects;
};

// 声明函数
int read_image_opencv(const cv::Mat& image, image_buffer_t* image_buffer);
int show_image(const image_buffer_t* image_buffer);
// cv::Mat process_image(const image_buffer_t* image_buffer);
cv::Mat processImage(rknn_app_context_t* rknn_app_ctx, const cv::Mat& image);

// 目标检测和3D定位函数
DetectionResult TargetFind(
    rknn_app_context_t* rknn_app_ctx,
    const cv::Mat& rgb_image,
    const cv::Mat& depth_image,
    const cv::Mat& camera_matrix,
    const std::vector<std::string>& target_classes,
    const std_msgs::msg::Header& header,
    tf2_ros::Buffer& tf_buffer,
    rclcpp::Logger logger
);

int read_image_opencv(const cv::Mat& image, image_buffer_t* image_buffer) 
{  
    
    // 使用 OpenCV 读取图像  
    if (image.empty()) 
    {  
        printf("error: read image fail\n");  
        return -1;  
    }  
  
    // 确定图像格式和通道数  
    int channels = image.channels();  
    image_buffer->format = (channels == 4) ? IMAGE_FORMAT_RGBA8888 :  
                           (channels == 1) ? IMAGE_FORMAT_GRAY8 :  
                                             IMAGE_FORMAT_RGB888;  
  
    // 设置图像宽度和高度  
    image_buffer->width = image.cols;  
    image_buffer->height = image.rows;  
  
    // 分配内存并复制图像数据  
    int size = image.total() * channels;  
    if (image_buffer->virt_addr != NULL) 
    {  
        // 如果 image->virt_addr 已经分配了内存，则复制数据到该内存  
        memcpy(image_buffer->virt_addr, image.data, size);  
    } 
    else 
    {  
        // 否则，分配新内存  
        image_buffer->virt_addr = (unsigned char *)malloc(size); 
        if (image_buffer->virt_addr == NULL) {  
            printf("error: memory allocation fail\n");  
            return -1;  
        }  
        memcpy(image_buffer->virt_addr, image.data, size);  
    }  
  
    // 如果图像是 4 通道（RGBA），但我们需要 3 通道（RGB），则可以选择性地转换  
    // 这里假设我们需要 RGB，所以如果是 RGBA，则去掉 A 通道  
    if (channels == 4) 
    {  
        cv::Mat rgb_img;  
        cv::cvtColor(image, rgb_img, cv::COLOR_RGBA2RGB);  
        memcpy(image_buffer->virt_addr, rgb_img.data, rgb_img.total() * 3);  
        // 更新大小（去掉 A 通道后的新大小）  
        size = rgb_img.total() * 3;  
    }  
  
    // 注意：这里我们没有释放 cv_img 的内存，因为 OpenCV 会自动管理它的内存。  
    // 但是，我们分配给 image->virt_addr 的内存需要在使用完后手动释放。  
  
    return 0;  
}

int show_image(const image_buffer_t* image_buffer) {  
    int width = image_buffer->width;  
    int height = image_buffer->height;  
    int channels = (image_buffer->format == IMAGE_FORMAT_RGB888) ? 3 :   
                    (image_buffer->format == IMAGE_FORMAT_GRAY8) ? 1 :   
                   4; // 根据image_buffer_t中的format字段确定通道数  
    void* data = image_buffer->virt_addr;  
  
    // 构造OpenCV Mat
    cv::Mat cv_img(height, width, CV_8UC(channels), data);  
  
    // 如果是RGB888格式，OpenCV显示需要转为BGR
    if (channels == 3 && image_buffer->format == IMAGE_FORMAT_RGB888) {
        cv::Mat bgr_img;
        cv::cvtColor(cv_img, bgr_img, cv::COLOR_RGB2BGR);
        cv::imshow("Detection Result", bgr_img);
    } else {
        cv::imshow("Detection Result", cv_img);
    }
    cv::waitKey(1); // 实时显示，1ms刷新

    return 0;
}

// cv::Mat process_image(const image_buffer_t* image_buffer){
//     int width = image_buffer->width;  
//     int height = image_buffer->height;  
//     int channels = (image_buffer->format == IMAGE_FORMAT_RGB888) ? 3 :   
//                     (image_buffer->format == IMAGE_FORMAT_GRAY8) ? 1 :   
//                    4; // 根据image_buffer_t中的format字段确定通道数  
//     void* data = image_buffer->virt_addr;  
  
//     // 构造OpenCV Mat
//     cv::Mat cv_img(height, width, CV_8UC(channels), data);  
  
//     // 如果是RGB888格式，OpenCV显示需要转为BGR
//     if (channels == 3 && image_buffer->format == IMAGE_FORMAT_RGB888) {
//         cv::Mat bgr_img;
//         cv::cvtColor(cv_img, bgr_img, cv::COLOR_RGB2BGR);
//         return bgr_img;
//     } else
//     {
//         return cv_img;
//     }
// }
// // 处理图像 
// cv::Mat processImage(rknn_app_context_t* rknn_app_ctx, const cv::Mat& image) 
// {
//     if (image.empty()) {
//         printf("input image is empty!\n");
//         return image;
//     }

//     int ret;
//     image_buffer_t src_image;
//     memset(&src_image, 0, sizeof(image_buffer_t));

//     ret = read_image_opencv(image, &src_image);

//     if (ret != 0) {  
//         printf("read image fail! ret=%d", ret);  
//         }  

//     object_detect_result_list od_results;
//     ret = inference_yolov8_model(rknn_app_ctx, &src_image, &od_results);
//     if (ret != 0) {  
//         printf("inference_yolov8_model fail! ret=%d\n", ret);  
//         if (src_image.virt_addr != NULL) {  
//             free(src_image.virt_addr);  
//             }  
//         } 

//     // 画框和概率
//     char text[256];
//     for (int i = 0; i < od_results.count; i++) {
//         object_detect_result *det_result = &(od_results.results[i]);
//         // printf("%s @ (%d %d %d %d) %.3f\n", coco_cls_to_name(det_result->cls_id),
//         //        det_result->box.left, det_result->box.top,
//         //        det_result->box.right, det_result->box.bottom,
//         //        det_result->prop);
//         int x1 = det_result->box.left;
//         int y1 = det_result->box.top;
//         int x2 = det_result->box.right;
//         int y2 = det_result->box.bottom;

//         // 这里的绘图现在是在 RGB 数据上进行
//         draw_rectangle(&src_image, x1, y1, x2 - x1, y2 - y1, COLOR_BLUE, 3);
//         sprintf(text, "%s %.1f%%", coco_cls_to_name(det_result->cls_id), det_result->prop * 100);
//         draw_text(&src_image, text, x1, y1 - 20, COLOR_RED, 10);
//     }
    
//     // 展示图片
//     // show_image(&src_image);
//     cv::Mat result_image = process_image(&src_image);
    
//     // 目标点处理
//     return result_image;
//     // 不需要手动free image.data
// }

// 计算目标点应朝向机器人的方向
geometry_msgs::msg::Quaternion calculate_orientation(
    const geometry_msgs::msg::Point& target,
    const geometry_msgs::msg::Point& robot,
    bool face_target = true)  // 添加选项参数
{
    double dx, dy;
    
    if (face_target) {
        // 机器人朝向目标点
        dx = target.x - robot.x;
        dy = target.y - robot.y;
    } else {
        // 目标点朝向机器人（当前实现）
        dx = robot.x - target.x;
        dy = robot.y - target.y;
    }
    
    double yaw = std::atan2(dy, dx);
    // double yaw = 0.0;
    
    tf2::Quaternion q;
    q.setRPY(0, 0, yaw);
    
    return tf2::toMsg(q);
}

float get_depth_value(const cv::Mat& depth_image, int x, int y) 
{
    std::cout<<"x: "<<x<<std::endl
        <<"y: "<<y<<std::endl;
    if (x < 0 || x >= depth_image.cols || y < 0 || y >= depth_image.rows) {
        return 0.0f;
    }
    
    // 假设深度图像是16UC1格式（毫米单位）
    return static_cast<float>(depth_image.at<uint16_t>(y, x)) / 1000.0f; // 转换为米
}

bool is_valid_depth(float depth) 
{
    return (depth > 0.1f) && (depth <= 10.0f);
}

float get_multi_point_depth(const cv::Mat& depth_image, const cv::Rect& bbox, rclcpp::Logger logger) 
{
    // 定义采样点网格
    const int grid_x = 10; // x方向点数
    const int grid_y = 10; // y方向点数
    std::vector<float> valid_depths;
    
    for (int i = 0; i < grid_x; i++) {
        for (int j = 0; j < grid_y; j++) {
            int x = bbox.x + (i * bbox.width) / (grid_x - 1);
            int y = bbox.y + (j * bbox.height) / (grid_y - 1);
            
            // 确保点在边界框内
            x = std::clamp(x, bbox.x, bbox.x + bbox.width - 1);
            y = std::clamp(y, bbox.y, bbox.y + bbox.height - 1);
            
            float depth = get_depth_value(depth_image, x, y);
            std::cout<<"深度值："<<depth<<std::endl;
            if (is_valid_depth(depth)) {
                std::cout<<"入队"<<std::endl;
                valid_depths.push_back(depth);
            }
        }
    }
    
    // 检查有效点数量
    if (valid_depths.empty()) {
        RCLCPP_WARN(logger, "多点采样无有效深度值");
        return 0.0f;
    }
    
    // 返回中值深度
    std::sort(valid_depths.begin(), valid_depths.end());
    return valid_depths[valid_depths.size()/2];
}

// 实现 TargetFind 函数
DetectionResult TargetFind(
    rknn_app_context_t* rknn_app_ctx,
    const cv::Mat& rgb_image,
    const cv::Mat& depth_image,
    const cv::Mat& camera_matrix,
    const std::vector<std::string>& target_classes,
    const std_msgs::msg::Header& header,
    tf2_ros::Buffer& tf_buffer,
    rclcpp::Logger logger)
{
    DetectionResult result;
    result.processed_image = rgb_image.clone(); // 初始化为原始图像副本
    
    // 1. 执行目标检测
    if (rgb_image.empty()) {
        RCLCPP_WARN(logger, "输入图像为空!");
        return result;
    }

    int ret;
    image_buffer_t src_image;
    memset(&src_image, 0, sizeof(image_buffer_t));

    ret = read_image_opencv(rgb_image, &src_image);
    if (ret != 0) {
        RCLCPP_ERROR(logger, "读取图像失败! ret=%d", ret);
        return result;
    }

    object_detect_result_list od_results;
    ret = inference_yolov8_model(rknn_app_ctx, &src_image, &od_results);
    if (ret != 0) {
        RCLCPP_ERROR(logger, "推理失败! ret=%d", ret);
        if (src_image.virt_addr != NULL) {
            free(src_image.virt_addr);
        }
        return result;
    }

    // 2. 处理检测结果
    for (int i = 0; i < od_results.count; i++) {
        object_detect_result* det_result = &(od_results.results[i]);
        std::string class_name = coco_cls_to_name(det_result->cls_id);
        
        // 检查是否是关注的目标类别
        if (std::find(target_classes.begin(), target_classes.end(), class_name) == target_classes.end()) {
            continue;
        }

        // 创建检测对象
        DetectedObject obj;
        obj.label = class_name;
        obj.confidence = det_result->prop;
        obj.bbox = cv::Rect(
            det_result->box.left,
            det_result->box.top,
            det_result->box.right - det_result->box.left,
            det_result->box.bottom - det_result->box.top
        );
        
        // 使用多点采样获取深度值
        float depth = get_multi_point_depth(depth_image, obj.bbox, logger);
        
        // 处理无效深度值
        if (depth <= 0.1f || depth > 10.0f) {
            RCLCPP_WARN(logger, "无效深度值: %.2f m", depth);
            continue;
        }
        
        // 计算边界框中心点
        int center_x = (det_result->box.left + det_result->box.right) / 2;
        int center_y = (det_result->box.top + det_result->box.bottom) / 2;
        
        // 3. 将2D像素坐标转换为3D相机坐标
        cv::Mat uv_point = (cv::Mat_<double>(3, 1) << center_x, center_y, 1.0);
        cv::Mat camera_point = camera_matrix.inv() * uv_point * depth;
        
        // 填充3D位置（相机坐标系）
        obj.position_3d.header = header;
        obj.position_3d.header.frame_id = "camera_color_optical_frame"; // 使用光学坐标系
        obj.position_3d.point.x = camera_point.at<double>(0);   // 0 y
        obj.position_3d.point.y = camera_point.at<double>(1);   // 1 z
        obj.position_3d.point.z = camera_point.at<double>(2);   // 2 x
        
       std::cout<<"相机坐标系： "<<std::endl
                <<"x: "<<obj.position_3d.point.x<<std::endl
                <<"y:"<<obj.position_3d.point.y<<std::endl
                <<"z"<<obj.position_3d.point.z<<std::endl;
        // 5. 转换到地图坐标系
        try {
    // 使用最新可用变换
            auto transform = tf_buffer.lookupTransform(
                "map", 
                obj.position_3d.header.frame_id,
                tf2::TimePointZero,  // 关键修改：使用最新变换
                tf2::durationFromSec(0.1));

            geometry_msgs::msg::PointStamped map_point;
            tf2::doTransform(obj.position_3d, map_point, transform);
            obj.position_3d = map_point;
            std::cout<<"地图坐标系： "<<std::endl
            <<"x: "<<obj.position_3d.point.x<<std::endl
            <<"y: "<<obj.position_3d.point.y<<std::endl
            <<"z: "<<obj.position_3d.point.z<<std::endl;
        } catch (tf2::TransformException &ex) {
            RCLCPP_WARN(logger, "坐标转换失败: %s", ex.what());
        }
        
        // 添加到结果列表
        result.detected_objects.push_back(obj);
        
        // 6. 在图像上绘制结果
        cv::rectangle(result.processed_image, obj.bbox, cv::Scalar(0, 255, 0), 2);
        std::string label_text = obj.label + ": " + std::to_string(obj.confidence * 100) + "%";
        cv::putText(result.processed_image, label_text, 
                    cv::Point(obj.bbox.x, obj.bbox.y - 10), 
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1);
    }

    // 释放资源
    if (src_image.virt_addr != NULL) {
        free(src_image.virt_addr);
    }

    return result;
}

